import pandas as pd
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.naive_bayes import MultinomialNB
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score, confusion_matrix
import jieba
import joblib

data = pd.read_csv('The Martian_reviews_cleaned.csv')

X = data['Review']
y = data['Sentiment']

X_segmented = X.apply(lambda x: ' '.join(jieba.cut(x)))

# 构建 TF-IDF 特征
vectorizer = TfidfVectorizer()
X_vectorized = vectorizer.fit_transform(X_segmented)


# 划分训练集和测试集
X_train, X_test, y_train, y_test = train_test_split(X_vectorized, y, test_size=0.8, random_state=42)

# 训练朴素贝叶斯模型
clf = MultinomialNB()
clf.fit(X_train, y_train)

# 保存模型
joblib.dump(clf, 'sentiment_classifier.pkl')
joblib.dump(vectorizer, 'tfidf_vectorizer.pkl')

# 评估模型性能
y_pred = clf.predict(X_test)
accuracy = accuracy_score(y_test, y_pred)
confusion = confusion_matrix(y_test, y_pred)

print(f'准确率: {accuracy:}')
# 加载保存的模型
clf = joblib.load('sentiment_classifier.pkl')
vectorizer = joblib.load('tfidf_vectorizer.pkl')

# 使用模型进行预测
new_text = input('请输入评论：')
new_text_seg = ' '.join(jieba.cut(new_text))
new_text_feature = vectorizer.transform([new_text_seg])
predicted_label = clf.predict(new_text_feature)[0]
print(f'新文本的情感预测结果为: {predicted_label}')